import random
import os

import yaml
from PIL import Image
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms as transforms

# define a training image loader that specifies transforms on images. See documentation for more details.
train_transformer = transforms.Compose([
    transforms.Resize(224),  # resize the image to 64x64 (remove if images are already 64x64)
    transforms.RandomHorizontalFlip(),  # randomly flip image horizontally
    transforms.ToTensor()])  # transform it into a torch tensor

# loader for evaluation, no horizontal flip
eval_transformer = transforms.Compose([
    transforms.Resize(224),  # resize the image to 64x64 (remove if images are already 64x64)
    transforms.ToTensor()])  # transform it into a torch tensor


class SIGNSDataset(Dataset):
    """
    A standard PyTorch definition of Dataset which defines the functions __len__ and __getitem__.
    """
    def __init__(self, data_dir, transform):
        """
        Store the filenames of the jpgs to use. Specifies transforms to apply on images.

        Args:
            data_dir: (string) directory containing the dataset
            transform: (torchvision.transforms) transformation to apply on image
        """
        self.filenames = os.listdir(data_dir)
        self.filenames = [os.path.join(data_dir, f) for f in self.filenames if f.endswith('.jpg')]

        self.labels = [int(os.path.split(filename)[-1][0]) for filename in self.filenames]
        self.transform = transform

    def __len__(self):
        # return size of dataset
        return len(self.filenames)

    def __getitem__(self, idx):
        """
        Fetch index idx image and labels from dataset. Perform transforms on image.

        Args:
            idx: (int) index in [0, 1, ..., size_of_dataset-1]

        Returns:
            image: (Tensor) transformed image
            label: (int) corresponding label of image
        """
        image = Image.open(self.filenames[idx])  # PIL image
        image = self.transform(image)
        return image, self.labels[idx]


def fetch_dataloader(types, config):
    """
    Fetches the DataLoader object for each type in types from data_dir.

    Args:
        types: (list) has one or more of 'train', 'val', 'test' depending on which data is required
        config: (yaml) hyperparameters

    Returns:
        data: (dict) contains the DataLoader object for each type in types
    """
    dataloaders = {}

    for split in ['train', 'val', 'test']:
        if split in types:
            path = os.path.join(config['data']['data_dir'], "{}_signs".format(split))

            # use the train_transformer if training data, else use eval_transformer without random flip
            if split == 'train':
                dl = DataLoader(
                    SIGNSDataset(path, train_transformer),
                    batch_size=config['data']['batch_size'],
                    shuffle=True,
                    num_workers=config['data']['num_workers'],
                    pin_memory=config['device']['gpu']
                )
            else:
                dl = DataLoader(
                    SIGNSDataset(path, eval_transformer),
                    batch_size=config['data']['batch_size'],
                    shuffle=False,
                    num_workers=config['data']['num_workers'],
                    pin_memory=config['device']['gpu']
                )

            dataloaders[split] = dl

    return dataloaders


if __name__ == '__main__':
    yaml_path = '../experiments/yaml_config/config.yaml'
    # print(yaml_path)
    with open(yaml_path, "r", encoding='utf-8') as f:
        config = yaml.safe_load(f)
    assert os.path.isfile(
        yaml_path), "No yaml configuration file found at {}".format(yaml_path)
    dataloaders = fetch_dataloader(['test'], config)
    test_dl = dataloaders['test']

    for data_batch, labels_batch in test_dl:
        print(data_batch.dtype, data_batch.size())
        print(labels_batch.dtype, labels_batch.size())
        break
